function hgmm = hgmm_train(X,y,cluster_mapping,num_cluster_components, num_class_components)

[classes, class_cluster_indices] = sort(unique(y));
clusters = sort(unique(cluster_mapping));
num_classes = size(classes, 1);
num_clusters = size(clusters, 1);
num_points = size(X, 1);

hgmm.cluster_model = cell(num_clusters, 1);
hgmm.cluster_prior = ones(num_clusters, 1);

hgmm.class_cluster = zeros(num_points, 1);

for i = 1 : num_clusters
    cluster = clusters(i);
    Xi = X(cluster_mapping(y) == cluster,:);
    cluster_size = size(Xi, 1);
    hgmm.cluster_prior(i) = cluster_size / num_points;
    fprintf('Fitting GMM to cluster %d ...\n', cluster);
    if (num_cluster_components == 1)
        hgmm.cluster_model{i} = gmm_fitml(Xi);
    else
        hgmm.cluster_model{i} = gmm_fitem(Xi, num_cluster_components);
    end
    hgmm.cluster_label(i) = cluster;
    hgmm.class_cluster(cluster_mapping(y) == cluster) = i;
end

hgmm.class_model = cell(num_classes, 1);
hgmm.class_prior = ones(num_classes, 1);

for i = 1 : num_classes
    class = classes(i);
    cluster = cluster_mapping(class_cluster_indices(i));
    Xi = X(y == class,:);
    class_size = size(Xi, 1);
    % hgmm.class_prior(i) = class_size / sum(cluster_mapping(y) == cluster);
    hgmm.class_prior(i) = class_size / num_points;
    fprintf('Fitting GMM to class %d ...\n', class);
    if (num_class_components == 1)
        hgmm.class_model{i} = gmm_fitml(Xi);
    else
        hgmm.class_model{i} = gmm_fitem(Xi, num_class_components);
    end
    hgmm.class_label(i) = class;
end

